package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"stock-analyzer/internal/database"
	"stock-analyzer/internal/model"
	"stock-analyzer/internal/service"
	"strconv"
	"strings"
	"time"

	"github.com/spf13/cobra"
)

var (
	dbPath      string
	startDate   string
	endDate     string
	maxDrawdown float64
	windowRules string
	outputFormat string
)

var rootCmd = &cobra.Command{
	Use:   "stock-analyzer",
	Short: "股票分析工具 - 基于回撤和滑动窗口涨跌幅筛选股票",
	Long: `股票分析工具，支持通过回撤范围和滑动窗口涨跌幅条件筛选股票。
示例: stock-analyzer --db stocks.db --start 2024-01-01 --end 2024-06-30 --max-drawdown 10 --windows "5:-3:lt,10:-5:lt,20:10:gt"`,
	Run: func(cmd *cobra.Command, args []string) {
		runAnalysis()
	},
}

func init() {
	rootCmd.PersistentFlags().StringVar(&dbPath, "db", "stocks.db", "SQLite数据库文件路径")
	rootCmd.PersistentFlags().StringVar(&startDate, "start", time.Now().AddDate(0, 0, -60).Format("2006-01-02"), "开始日期 (格式: YYYY-MM-DD)")
	rootCmd.PersistentFlags().StringVar(&endDate, "end", time.Now().Format("2006-01-02"), "结束日期 (格式: YYYY-MM-DD)")
	rootCmd.PersistentFlags().Float64Var(&maxDrawdown, "max-drawdown", 10.0, "最大允许回撤百分比")
	rootCmd.PersistentFlags().StringVar(&windowRules, "windows", "", "窗口涨跌幅规则，格式: '窗口:涨跌幅:类型,窗口:涨跌幅:类型'")
	
	// 添加输出格式选项
	rootCmd.PersistentFlags().StringVarP(&outputFormat, "output", "o", "text", "输出格式 (text, json, csv)")
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func runAnalysis() {
	// 初始化数据库
	db, err := database.NewDB(dbPath)
	if err != nil {
		log.Fatalf("Failed to initialize database: %v", err)
	}
	defer db.Close()

	// 解析日期
	start, err := time.Parse("2006-01-02", startDate)
	if err != nil {
		log.Fatalf("无效的开始日期格式: %v", err)
	}

	end, err := time.Parse("2006-01-02", endDate)
	if err != nil {
		log.Fatalf("无效的结束日期格式: %v", err)
	}

	// 解析窗口规则
	windowChanges, err := parseWindowRules(windowRules)
	if err != nil {
		log.Fatalf("Failed to parse window rules: %v", err)
	}

	// 创建筛选条件
	condition := model.FilterCondition{
		StartDate:     start,
		EndDate:       end,
		MaxDrawdown:   &maxDrawdown,
		WindowChanges: windowChanges,
	}

	// 创建筛选服务
	service := service.NewFilterService(db)

	// 执行筛选
	results, err := service.FilterStocks(condition)
	if err != nil {
		log.Fatalf("Failed to filter stocks: %v", err)
	}

	// 输出结果
	printResults(results, condition)
}

func parseWindowRules(rulesStr string) ([]model.WindowChangeRule, error) {
	if rulesStr == "" {
		return nil, nil
	}

	var rules []model.WindowChangeRule
	ruleParts := strings.Split(rulesStr, ",")

	for _, part := range ruleParts {
		fields := strings.Split(part, ":")
		if len(fields) != 3 {
			return nil, fmt.Errorf("invalid window rule format: %s", part)
		}

		window, err := strconv.Atoi(fields[0])
		if err != nil {
			return nil, fmt.Errorf("invalid window size: %s", fields[0])
		}

		change, err := strconv.ParseFloat(fields[1], 64)
		if err != nil {
			return nil, fmt.Errorf("invalid change value: %s", fields[1])
		}

		ruleType := fields[2]
		if !isValidRuleType(ruleType) {
			return nil, fmt.Errorf("invalid rule type: %s (must be gt, lt, ge, le)", ruleType)
		}

		rules = append(rules, model.WindowChangeRule{
			Window: window,
			Change: change,
			Type:   ruleType,
		})
	}

	return rules, nil
}

func isValidRuleType(ruleType string) bool {
	switch ruleType {
	case "gt", "lt", "ge", "le":
		return true
	default:
		return false
	}
}

func printResults(results []model.StockResult, condition model.FilterCondition) {
	// 使用全局变量outputFormat

	// 如果是CSV格式，直接输出CSV
	if outputFormat == "csv" {
		printCSVResults(results, condition)
		return
	}

	// 文本格式输出
	fmt.Printf("筛选条件:\n")
	fmt.Printf("  时间范围: %s 至 %s\n", condition.StartDate.Format("2006-01-02"), condition.EndDate.Format("2006-01-02"))
	if condition.MaxDrawdown != nil {
		fmt.Printf("  最大回撤: %.2f%%\n", *condition.MaxDrawdown)
	}
	for _, rule := range condition.WindowChanges {
		fmt.Printf("  %d天窗口: %s %.2f%%\n", rule.Window, getRuleTypeDesc(rule.Type), rule.Change)
	}
	fmt.Printf("\n")

	if len(results) == 0 {
		fmt.Println("没有找到符合条件的股票")
		return
	}

	fmt.Printf("找到 %d 只符合条件的股票:\n", len(results))
	fmt.Printf("%-10s %-8s", "股票代码", "回撤")
	for _, rule := range condition.WindowChanges {
		fmt.Printf(" %-8s", fmt.Sprintf("%d天", rule.Window))
	}
	fmt.Printf("\n")

	for _, result := range results {
		fmt.Printf("%-12s %-8.2f%%", result.StockCode, result.MaxDrawdown)
		for _, rule := range condition.WindowChanges {
			for _, wc := range result.WindowChanges {
				if wc.Window == rule.Window {
					fmt.Printf(" %-8.2f%%", wc.Change)
					break
				}
			}
		}
		fmt.Printf("\n")
	}

	// 输出JSON格式结果（可选）
	if outputFormat == "json" && len(results) > 0 {
		if jsonOutput, _ := json.MarshalIndent(results, "", "  "); len(results) > 0 {
			fmt.Printf("\nJSON格式结果:\n%s\n", string(jsonOutput))
		}
	}
}

func getRuleTypeDesc(ruleType string) string {
	switch ruleType {
	case "gt":
		return "大于"
	case "lt":
		return "小于"
	case "ge":
		return "大于等于"
	case "le":
		return "小于等于"
	default:
		return ruleType
	}
}

// 输出CSV格式的结果
func printCSVResults(results []model.StockResult, condition model.FilterCondition) {
	// 如果没有结果，只输出头部
	if len(results) == 0 {
		fmt.Println("股票代码,最大回撤")
		return
	}

	// 构建CSV头部
	header := "股票代码,最大回撤"
	for _, rule := range condition.WindowChanges {
		header += fmt.Sprintf(",%d天涨跌幅", rule.Window)
	}
	fmt.Println(header)

	// 输出每一行数据
	for _, result := range results {
		line := fmt.Sprintf("%s,%.2f", result.StockCode, result.MaxDrawdown)
		
		// 添加每个窗口的涨跌幅
		for _, rule := range condition.WindowChanges {
			found := false
			for _, wc := range result.WindowChanges {
				if wc.Window == rule.Window {
					line += fmt.Sprintf(",%.2f", wc.Change)
					found = true
					break
				}
			}
			if !found {
				line += ","
			}
		}
		
		fmt.Println(line)
	}
}
